
package org.apache.solr.search;

import java.io.IOException;
import java.util.Map;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.ValueSourceScorer;
import org.apache.lucene.search.IndexSearcher;
import org.apache.solr.search.function.ValueSourceRangeFilter;

// This class works as either a normal constant score query, or as a PostFilter using a collector
public class FunctionRangeQuery extends SolrConstantScoreQuery implements PostFilter {

    final ValueSourceRangeFilter rangeFilt;

    public FunctionRangeQuery(ValueSourceRangeFilter filter) {
        super(filter);
        this.rangeFilt = filter;
    }

    @Override
    public DelegatingCollector getFilterCollector(IndexSearcher searcher) {
        Map fcontext = ValueSource.newContext(searcher);
        return new FunctionRangeCollector(fcontext);
    }

    class FunctionRangeCollector extends DelegatingCollector {

        final Map fcontext;
        ValueSourceScorer scorer;
        int maxdoc;

        public FunctionRangeCollector(Map fcontext) {
            this.fcontext = fcontext;
        }

        @Override
        public void collect(int doc) throws IOException {
            if (doc < maxdoc && scorer.matches(doc)) {
                delegate.collect(doc);
            }
        }

        @Override
        public void setNextReader(AtomicReaderContext context) throws IOException {

            maxdoc = context.reader().maxDoc();
            FunctionValues dv = rangeFilt.getValueSource().getValues(fcontext, context);
            scorer = dv.getRangeScorer(context.reader(), rangeFilt.getLowerVal(), rangeFilt.getUpperVal(), rangeFilt.isIncludeLower(), rangeFilt.isIncludeUpper());
            super.setNextReader(context);
        }
    }
}
